import numpy as np
from . import LocalSearchOperator, register_operator


class TwoOptOperator(LocalSearchOperator):
    """通用的2-opt局部搜索算子"""

    def __init__(self):
        self.name = "two_opt"
        self.supported_problems = ["tsp"]

    def generate_neighborhood(self, solution, problem_instance):
        """生成2-opt邻域"""
        neighborhood = []
        n = len(solution)

        for i in range(1, n - 2):
            for j in range(i + 1, n):
                if j - i == 1:
                    continue
                new_solution = solution.copy()
                new_solution[i:j + 1] = new_solution[i:j + 1][::-1]
                neighborhood.append(new_solution)

        return neighborhood

    def __call__(self, solution, problem_instance, intensity):
        """应用2-opt局部搜索"""
        if not self.is_applicable(problem_instance):
            return solution, 0

        best_solution = solution.copy()
        best_cost = problem_instance.evaluate(best_solution)
        improved = True
        iterations = 0

        while improved and iterations < intensity:
            improved = False
            neighborhood = self.generate_neighborhood(best_solution, problem_instance)

            for neighbor in neighborhood:
                neighbor_cost = problem_instance.evaluate(neighbor)

                if neighbor_cost < best_cost:
                    best_solution = neighbor
                    best_cost = neighbor_cost
                    improved = True
                    break

            iterations += 1

        improvement = problem_instance.evaluate(solution) - best_cost
        return best_solution, improvement


# 注册算子
register_operator("two_opt", TwoOptOperator())


